import os
import mne
import torch
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch.utils.data.dataset as Dataset
# 设置环境变量避免库冲突
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# 设置seed
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)

# cuda
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class EEGDataset(Dataset.Dataset):
    def __init__(self, Data, Label):
        self.Data = Data
        self.Label = Label
        print("Data shape: ", self.Data.shape)
        print("Label shape: ", self.Label.shape)

    def __len__(self):
        return len(self.Data)

    def __getitem__(self, index):
        data = torch.Tensor(self.Data[index]).unsqueeze(-1).to(device)
        label = torch.Tensor(self.Label[index]).to(device)
        # 打印数据和标签的形状
        # print(f"data shape: {data.shape}, label shape: {label.shape}")
        return data, label


# 读取指定路径下的文件数据
def load_data_BCICIV_2b_gdf(root_dir="./BCICIV_2b_gdf/train_data"):
    # 打印loading信息
    print("Loading BCICIV_2b data from: ", root_dir)
    # 存储加窗后扩展的数据和标签
    # 窗口为2s(500个采样点)，重叠为200ms(50个采样点)，一共取11个窗口
    expend_data = []
    expend_label = []
    # 打开文件夹
    files = os.listdir(root_dir)
    for file in tqdm(files):
        # print("Loading file: ", file)
        file_path = os.path.join(root_dir, file)
        # 读取gdf文件 默认低通0，高通为250hz/2
        raw_gdf = mne.io.read_raw_gdf(file_path, stim_channel='auto', verbose='ERROR',
                                      exclude=(["EOG:ch01", "EOG:ch02", "EOG:ch03", "EEG:Cz", "EEG:C4"]))

        raw_gdf.load_data(verbose='ERROR')
        data = raw_gdf.get_data()

        for i_chan in range(data.shape[0]):  # 遍历channels
            # 将数组中的所有值设置为nan，然后将这些NaN值替换为该数组的均值。
            this_chan = data[i_chan]
            data[i_chan] = np.where(
                this_chan == np.min(this_chan), np.nan, this_chan
            )
            mask = np.isnan(data[i_chan])
            chan_mean = np.nanmean(data[i_chan])
            data[i_chan, mask] = chan_mean
        # 获取事件时间位置，返回事件和事件下标
        events, events_id = mne.events_from_annotations(raw_gdf, verbose='ERROR')
        # 如果没有769和770事件，跳过
        # print(events_id)
        if '769' not in events_id or '770' not in events_id:
            print('WARNING: No 769 or 770 event in file:', file)
            continue
        # print('Number of events:', len(events))
        # print(events_id)
        # 利用mne.io.RawArray类重新创建Raw对象，已经没有nan数据了
        raw_gdf = mne.io.RawArray(data, raw_gdf.info, verbose="ERROR")
        # 选择范围为Cue后 1s - 5s 的数据
        tmin, tmax = 1., 5.
        # 二分类 MI 对应的 events_id
        MI_event_id = dict({'769': events_id['769'], '770': events_id['770']})

        epochs = mne.Epochs(raw_gdf, events, MI_event_id, tmin, tmax, proj=True, baseline=None, preload=True, verbose='ERROR')
        # print(epochs)
        # 将事件标签转换为0和1
        label_transfer = {events_id['769']: 0, events_id['770']: 1}
        # 切片，获取 events 的最后一列
        labels = epochs.events[:, -1]
        # 将labels转化为0和1
        labels = np.array([label_transfer[label] for label in labels])
        # print(labels)
        # Get all epochs as a 3D array.
        data = epochs.get_data()
        # 窗口为2s(500个采样点)，重叠为200ms(50个采样点)，一共取11个窗口
        for i in range(11):
            # 从i*50开始取500个采样点
            expend_data.append(data[:, :, i * 50:i * 50 + 500])
            expend_label.append(labels)
    # 按照0维度合并，单通道数据进行squeeze,标签增加一个维度
    expend_data = np.concatenate(expend_data, axis=0).squeeze(1)
    expend_label = np.concatenate(expend_label, axis=0).reshape(-1, 1)
    # print(expend_data.shape, expend_label.shape)
    print("取得样本数：", expend_data.shape[0])
    return expend_data, expend_label


def check_data_info(file_path, img_save=False, img_dir='./img'):
    # 读取gdf文件，exclude参数排除三个EOG通道
    raw_gdf = mne.io.read_raw_gdf(file_path, stim_channel="auto", verbose='ERROR',
                                  exclude=(["EOG:ch01", "EOG:ch02", "EOG:ch03"]))
    print(raw_gdf.info)
    raw_gdf.load_data()
    data = raw_gdf.get_data()

    for i_chan in range(data.shape[0]):  # 遍历 channel
        # 将数组中的所有值设置为nan，然后将这些NaN值替换为该数组的均值。
        this_chan = data[i_chan]
        data[i_chan] = np.where(
            this_chan == np.min(this_chan), np.nan, this_chan
        )
        mask = np.isnan(data[i_chan])
        chan_mean = np.nanmean(data[i_chan])
        data[i_chan, mask] = chan_mean

    # 获取事件时间位置，返回事件和事件下标
    events, events_id = mne.events_from_annotations(raw_gdf)
    # 主要需要注意的事件即MI数据：769, 770, 771, 772 分别为左手、右手、脚、舌头
    print('Number of events:', len(events))
    print(events_id)
    print(events.shape)

    # 利用mne.io.RawArray类重新创建Raw对象，已经没有nan数据了
    raw_gdf = mne.io.RawArray(data, raw_gdf.info, verbose="ERROR")
    print(raw_gdf.info)
    # 画出EEG通道图
    raw_gdf.plot()
    # 存储图像
    if img_save:
        if not os.path.exists(img_dir):
            os.makedirs(img_dir)
        plt.savefig(os.path.join(img_dir, 'EEG_channel.png'))
    plt.show()


# 打印项目使用的库版本
def print_version():
    print("numpy:", np.__version__)
    print("torch:", torch.__version__)
    print("mne:", mne.__version__)
    print("matplotlib:", plt.matplotlib.__version__)
    print("cuda:", torch.version.cuda)


